
import torchvision
import os
from torch.utils import data
import parameters as p
import cv2
from einops import rearrange
import torch



transforms = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Resize(32,
                                  antialias=None,
                                  interpolation=torchvision.transforms.InterpolationMode.BICUBIC)
])



def imshow(img:torch.Tensor):
    img = rearrange(img.numpy(),"b c h w ->  (b h) w c")
    cv2.imshow("s",img)
    cv2.waitKey(0)

def get_dataloader():
    
    dataset = torchvision.datasets.MNIST(root=p.DATASET_PATH,download=False,transform=transforms)
    val_dataset = torchvision.datasets.MNIST(root=p.DATASET_PATH,transform=transforms,train=False)
    rs = data.RandomSampler(val_dataset)
    return data.DataLoader(dataset,batch_size=p.BATCH_SIZE,num_workers=4,pin_memory=True,drop_last=True,persistent_workers=True),data.DataLoader(val_dataset,sampler=rs,batch_size=1)
# if __name__ == '__main__':
#     # for x,y in dataloader:
#     #     imshow(x)
#     #     print(x.shape)

#     #     print(y)
        
#     #     print(to_one_hot(y))
#     #     break
#     cv2.namedWindow('s',cv2.WINDOW_AUTOSIZE)